from pymilvus import MilvusClient

milvus_client = MilvusClient("http://localhost:19530")


def delete_data_by_expression(collection_name, pks, filter_expr):
    print(f"Start to delete by expression {filter_expr} in collection {collection_name}")

    try:
        # Delete data based on expression
        milvus_client.delete(collection_name, pks=pks, filter=filter_expr)
        milvus_client.flush(collection_name)
        print("Deletion successful.")
    except Exception as e:
        print(f"Failed to delete data: {e}")


def query_data_by_expression(collection_name, filter_expr, output_fields, limit=10):
    print(f"Start to query by expression {filter_expr} in collection {collection_name}")

    try:
        # Delete data based on expression
        query_result = milvus_client.query(collection_name, filter=filter_expr, output_fields=output_fields)

        query_result_limit = query_result[:limit]

        print("Query successful.")
    except Exception as e:
        print(f"Failed to Query data: {e}")

    return query_result_limit


def delete_data_by_user_id(user_id='1000000'):
    collection_name = "multimodal_search_H14"
    filter_expression = f"user_id == {user_id}"
    limit = 2000
    output_fields = ["id"]
    res = query_data_by_expression(collection_name, filter_expression, output_fields, limit)

    res_array = [item['id'] for item in res]

    delete_data_by_expression(collection_name, pks=res_array, filter_expr='')


def delete_data_by_id(ids):
    collection_name = "multimodal_search_H14"

    delete_data_by_expression(collection_name, pks=ids, filter_expr='')


def query_data():
    collection_name = "multimodal_search_H14"
    filter_expression = "user_id == '1000000'"
    limit = 1
    output_fields = ["id", "user_id", "image_path"]

    res = query_data_by_expression(collection_name, filter_expression, output_fields, limit)

    print(res)


def generate_string_array(start, end):
    """
    生成指定范围内的字符串数组

    Parameters:
    - start: 范围起始值
    - end: 范围结束值（不包含）

    Returns:
    - 字符串数组
    """
    return [str(i) for i in range(start, end + 1)]


if __name__ == "__main__":
    # delete_data_by_user_id('1000000')

    # 示例：生成范围字符串数组
    result_array = generate_string_array(1100122, 1100510)

    # 打印结果数组
    print(result_array)

    delete_data_by_id(result_array)
